!-------------------------------------------------------------LICENSE--------------------------------------------------------------!
!                                                                                                                                  !
!The MAP code is written in Fortran language for magnetohydrodynamics (MHD) calculation with the adaptive mesh refinement (AMR)    !
!and Message Passing Interface (MPI) parallelization.                                                                              !
!                                                                                                                                  !
!Copyright (C) 2012                                                                                                                !
!Ronglin Jiang                                                                                                                     !
!rljiang@ssc.net.cn                                                                                                                !
!585 Guoshoujing Road. Pudong, Shanghai, P.R.C. 201203                                                                             !
!                                                                                                                                  !
!This program is free software; you can redistribute it and/or modify it under the terms of the GNU General Public License         !
!as published by the Free Software Foundation; either version 2 of the License, or (at your option) any later version.             !
!                                                                                                                                  !
!This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of    !
!MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License for more details.                        !
!                                                                                                                                  !
!You should have received a copy of the GNU General Public License along with this program; if not, write to the Free Software     !
!Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.                                                   !
!                                                                                                                                  !
!-------------------------------------------------------------LICENSE--------------------------------------------------------------!

!==================================================================================================================================|
subroutine load_balance (base)
!==================================================================================================================================|

   use parameters
   use block_define
   implicit none

   type(block), pointer, intent(inout) :: base

   integer(4) :: i, j, i_block, i_cpu, i_hilbert, i_block_vip, i_request, i_level, n_var
   integer(4) :: average_blocks, blocks_p1_cpu_number, old_start, old_end, start_flag, send_start, send_end, n_request
   integer(4) :: new_start, new_end, n_send, n_recv, pos_x, pos_y
   integer(4), dimension(cpu_number) :: blocks_front, blocks_vip_front
   integer(4), dimension(cpu_number) :: hilbert_start_new, hilbert_end_new
   integer(4), dimension(n_blocks_info_send, cpu_number) :: blocks_info_send
   integer(4), dimension(n_blocks_info_send, cpu_number, cpu_number) :: blocks_info_sendrecv

   integer(4), dimension(:), allocatable :: request
   integer(4), dimension(:, :), allocatable :: request_status
   integer(4), dimension(:, :, :), allocatable :: info_int_send, info_int_recv
   real(8), dimension(:, :, :, :), allocatable :: info_var_send, info_var_recv

   type(block), pointer :: current, last

!----------------------------------------------------------------------------------------------------------------------------------|
!  flag the vip blocks in blocks_vip_flag and blocks_vip_flag_all
!  calculation the total number of the blocks_vip
!----------------------------------------------------------------------------------------------------------------------------------|
   if (allocated (blocks_vip_flag)) deallocate (blocks_vip_flag)
   if (allocated (blocks_vip_flag_all)) deallocate (blocks_vip_flag_all)
   allocate (blocks_vip_flag(n_blocks(cpu_id_p1)), blocks_vip_flag_all(n_blocks_all))

   blocks_vip_flag = 0
   blocks_vip_flag_all = 0

   i_block = 0
   n_blocks_vip(cpu_id_p1) = 0
   current => base%next
   do while (current%lv .ne. 0)
      i_block = i_block + 1
      if (current%nc .eq. 0) then
         n_blocks_vip(cpu_id_p1) = n_blocks_vip(cpu_id_p1) + 1
         blocks_vip_flag(i_block) = 1
      endif
      current => current%next
   enddo

   blocks_front = 0
   do i_cpu = 1, cpu_number - 1
      blocks_front(i_cpu + 1) = blocks_front(i_cpu) + n_blocks(i_cpu)
   enddo

   call mpi_allgatherv (blocks_vip_flag, n_blocks(cpu_id_p1), mpi_integer1,                                                        &
      blocks_vip_flag_all, n_blocks, blocks_front, mpi_integer1, mpi_comm_world, error)

!----------------------------------------------------------------------------------------------------------------------------------|
!  update the n_blocks_vip
!----------------------------------------------------------------------------------------------------------------------------------|
   call mpi_allreduce (n_blocks_vip(cpu_id_p1), n_blocks_vip_all, 1, mpi_integer4, mpi_sum, mpi_comm_world, error)
   average_blocks = n_blocks_vip(cpu_id_p1)
   call mpi_allgather (average_blocks, 1, mpi_integer4, n_blocks_vip, 1, mpi_integer4, mpi_comm_world, error)

!----------------------------------------------------------------------------------------------------------------------------------|
!  update the value of hilbert_start, hilbert_end and the hilbert number of every block
!----------------------------------------------------------------------------------------------------------------------------------|
   old_start = blocks_front(cpu_id_p1) + 1
   old_end = blocks_front(cpu_id_p1) + n_blocks(cpu_id_p1)

   call mpi_allgather (old_start, 1, mpi_integer4, hilbert_start, 1, mpi_integer4, mpi_comm_world, error)
   call mpi_allgather (old_end, 1, mpi_integer4, hilbert_end, 1, mpi_integer4, mpi_comm_world, error)

   i_hilbert = 0
   current => base%framework_next
   do while (current%lv .ne. 0)
      i_hilbert = i_hilbert + 1
      current%hilbert_number = i_hilbert
      current => current%framework_next
   enddo

!----------------------------------------------------------------------------------------------------------------------------------|
!  repartitioning
!----------------------------------------------------------------------------------------------------------------------------------|
   average_blocks = n_blocks_vip_all / cpu_number

   blocks_p1_cpu_number = mod(n_blocks_vip_all, cpu_number)
   if (cpu_id .lt. blocks_p1_cpu_number) average_blocks = average_blocks + 1

   call mpi_allgather (average_blocks, 1, mpi_integer4, n_blocks_vip, 1, mpi_integer4, mpi_comm_world, error)

   blocks_vip_front = 0
   do i_cpu = 2, cpu_number
      blocks_vip_front(i_cpu) = blocks_vip_front(i_cpu - 1) + n_blocks_vip(i_cpu - 1)
   enddo

   new_start = n_blocks_all + 1        ! for the situation that number of cpu > number of blocks
   new_end = n_blocks_all + 1
   i_block_vip = 0
   start_flag = 1
   do i_block = 1, n_blocks_all
      if (i_block_vip .eq. blocks_vip_front(cpu_id_p1) .and. start_flag .eq. 1) then
         new_start = i_block
         start_flag = 0
      endif
      if (blocks_vip_flag_all(i_block) .eq. 1) i_block_vip = i_block_vip + 1
      if (i_block_vip .eq. blocks_vip_front(cpu_id_p1) + n_blocks_vip(cpu_id_p1)) then
         new_end = i_block
         exit
      endif
   enddo

   call mpi_allgather (new_start, 1, mpi_integer4, hilbert_start_new, 1, mpi_integer4, mpi_comm_world, error)
   call mpi_allgather (new_end, 1, mpi_integer4, hilbert_end_new, 1, mpi_integer4, mpi_comm_world, error)

!----------------------------------------------------------------------------------------------------------------------------------|
!  collect the send information
!----------------------------------------------------------------------------------------------------------------------------------|
   blocks_info_send = 0
   do i_cpu = 1, cpu_number
      if (i_cpu .eq. cpu_id_p1) cycle
      send_start = max(hilbert_start(cpu_id_p1), hilbert_start_new(i_cpu))
      send_end = min(hilbert_end(cpu_id_p1), hilbert_end_new(i_cpu))
      if (send_end .ge. send_start) then
         blocks_info_send (1, i_cpu) = send_end - send_start + 1
         blocks_info_send (2, i_cpu) = send_start
         blocks_info_send (3, i_cpu) = send_end
      endif
   enddo

   call mpi_allgather (blocks_info_send, size (blocks_info_send), mpi_integer4,                                                    &
      blocks_info_sendrecv, size (blocks_info_send), mpi_integer4, mpi_comm_world, error)

!----------------------------------------------------------------------------------------------------------------------------------|
!  mpi status initialization
!----------------------------------------------------------------------------------------------------------------------------------|
   n_request = count (blocks_info_sendrecv(1, :, cpu_id_p1) .gt. 0) + count (blocks_info_sendrecv(1, cpu_id_p1, :) .gt. 0)
   allocate (request(n_request * 2), request_status(mpi_status_size, n_request * 2))

!----------------------------------------------------------------------------------------------------------------------------------|
!  update the cpu_id of every block
!----------------------------------------------------------------------------------------------------------------------------------|
   i_cpu = 0
   current => base%framework_next
   do while (current%lv .ne. 0)
      if (current%hilbert_number .ge. hilbert_start_new(i_cpu + 1) .and.                                                           &
          current%hilbert_number .le. hilbert_end_new(i_cpu + 1)) then
         current%cpu_id = i_cpu
      else
         i_cpu = i_cpu + 1
         cycle
      endif
      current => current%framework_next
   enddo

!----------------------------------------------------------------------------------------------------------------------------------|
!  send blocks
!----------------------------------------------------------------------------------------------------------------------------------|
   n_var = 12
   if (ini_value_save_flag .ne. 0) n_var = 21
   n_send = sum (blocks_info_sendrecv(1, :, cpu_id_p1))
   allocate (info_int_send(n_dims, n_levels, n_send), info_var_send(n_var, nx_block, ny_block, n_send))

   i_hilbert = 0
   do i_cpu = 0, cpu_number - 1
      if (i_cpu .eq. cpu_id) cycle
      n_send = blocks_info_sendrecv(1, i_cpu + 1, cpu_id_p1)
      if (n_send .eq. 0) cycle

      current => base%next
      do while (current%lv .ne. 0)
         if (current%hilbert_number .ge. blocks_info_sendrecv(2, i_cpu + 1, cpu_id_p1) .and.                                       &
             current%hilbert_number .le. blocks_info_sendrecv(3, i_cpu + 1, cpu_id_p1)) then
            i_hilbert = i_hilbert + 1

            n_blocks(cpu_id_p1) = n_blocks(cpu_id_p1) - 1

            info_int_send(:, :, i_hilbert) = current%id

            if (ini_value_save_flag .ne. 0) then
               do j = 1, ny_block
               do i = 1, nx_block
                  info_var_send(1, i, j, i_hilbert) = current%ro(i, j)
                  info_var_send(2, i, j, i_hilbert) = current%mx(i, j)
                  info_var_send(3, i, j, i_hilbert) = current%my(i, j)
                  info_var_send(4, i, j, i_hilbert) = current%mz(i, j)
                  info_var_send(5, i, j, i_hilbert) = current%bx(i, j)
                  info_var_send(6, i, j, i_hilbert) = current%by(i, j)
                  info_var_send(7, i, j, i_hilbert) = current%bz(i, j)
                  info_var_send(8, i, j, i_hilbert) = current%en(i, j)
                  info_var_send(9, i, j, i_hilbert) = current%az(i, j)
                  info_var_send(10, i, j, i_hilbert) = current%gx(i, j)
                  info_var_send(11, i, j, i_hilbert) = current%gy(i, j)
                  info_var_send(12, i, j, i_hilbert) = current%po(i, j)
                  info_var_send(13, i, j, i_hilbert) = current%ro0(i, j)
                  info_var_send(14, i, j, i_hilbert) = current%mx0(i, j)
                  info_var_send(15, i, j, i_hilbert) = current%my0(i, j)
                  info_var_send(16, i, j, i_hilbert) = current%mz0(i, j)
                  info_var_send(17, i, j, i_hilbert) = current%bx0(i, j)
                  info_var_send(18, i, j, i_hilbert) = current%by0(i, j)
                  info_var_send(19, i, j, i_hilbert) = current%bz0(i, j)
                  info_var_send(20, i, j, i_hilbert) = current%en0(i, j)
                  info_var_send(21, i, j, i_hilbert) = current%az0(i, j)
               enddo
               enddo
            else
               do j = 1, ny_block
               do i = 1, nx_block
                  info_var_send(1, i, j, i_hilbert) = current%ro(i, j)
                  info_var_send(2, i, j, i_hilbert) = current%mx(i, j)
                  info_var_send(3, i, j, i_hilbert) = current%my(i, j)
                  info_var_send(4, i, j, i_hilbert) = current%mz(i, j)
                  info_var_send(5, i, j, i_hilbert) = current%bx(i, j)
                  info_var_send(6, i, j, i_hilbert) = current%by(i, j)
                  info_var_send(7, i, j, i_hilbert) = current%bz(i, j)
                  info_var_send(8, i, j, i_hilbert) = current%en(i, j)
                  info_var_send(9, i, j, i_hilbert) = current%az(i, j)
                  info_var_send(10, i, j, i_hilbert) = current%gx(i, j)
                  info_var_send(11, i, j, i_hilbert) = current%gy(i, j)
                  info_var_send(12, i, j, i_hilbert) = current%po(i, j)
               enddo
               enddo
            endif

!            deallocate (current%ro, current%mx, current%my, current%mz, current%bx, current%by, current%bz, current%en,            &
!                        current%gx, current%gy, current%po, current%az)
!            current%var_allocate_flag = .false.
!            if (ini_value_save_flag .ne. 0) then
!               deallocate (current%ro0, current%mx0, current%my0, current%mz0, current%bx0, current%by0, current%bz0, current%en0, &
!                  current%az0)
!            endif
         endif
         current => current%next
      enddo
   enddo

   i_hilbert = 1
   i_request = 0
   do i_cpu = 0, cpu_number - 1
      if (i_cpu .eq. cpu_id) cycle
      n_send = blocks_info_sendrecv(1, i_cpu + 1, cpu_id_p1)
      if (n_send .eq. 0) cycle

      i_request = i_request + 1
      call mpi_isend (info_int_send(:, :, i_hilbert), n_dims * n_levels * n_send, mpi_integer4, i_cpu, tag_int,                    &
         mpi_comm_world, request(i_request), error)
      i_request = i_request + 1
      call mpi_isend (info_var_send(:, :, :, i_hilbert), n_var * nx_block * ny_block * n_send, mpi_real8, i_cpu, tag_var,          &
         mpi_comm_world, request(i_request), error)
      i_hilbert = i_hilbert + n_send
   enddo

!----------------------------------------------------------------------------------------------------------------------------------|
!  receive blocks
!----------------------------------------------------------------------------------------------------------------------------------|
   n_recv = sum (blocks_info_sendrecv(1, cpu_id_p1, :))
   allocate (info_int_recv(n_dims, n_levels, n_recv), info_var_recv(n_var, nx_block, ny_block, n_recv))

   i_hilbert = 1
   do i_cpu = 0, cpu_number - 1
      if (i_cpu .eq. cpu_id) cycle
      n_recv = blocks_info_sendrecv(1, cpu_id_p1, i_cpu + 1)
      if (n_recv .eq. 0) cycle

      i_request = i_request + 1
      call mpi_irecv (info_int_recv(:, :, i_hilbert), n_dims * n_levels * n_recv, mpi_integer4, i_cpu, tag_int,                    &
         mpi_comm_world, request(i_request), error)
      i_request = i_request + 1
      call mpi_irecv (info_var_recv(:, :, :, i_hilbert), n_var * nx_block * ny_block * n_recv, mpi_real8, i_cpu, tag_var,          &
         mpi_comm_world, request(i_request), error)
      i_hilbert = i_hilbert + n_recv
   enddo

   call mpi_waitall (size (request), request, request_status, error)

   n_recv = sum (blocks_info_sendrecv(1, cpu_id_p1, :))
   do i_hilbert = 1, n_recv
      i_level = 1
      current => base
      do while (info_int_recv(1, i_level, i_hilbert) .ne. 0)
         pos_x = info_int_recv(1, i_level, i_hilbert)
         pos_y = info_int_recv(2, i_level, i_hilbert)
         current => current%child(pos_x, pos_y)%pblock
         i_level = i_level + 1
         if (i_level .gt. n_levels) exit
      enddo

      n_blocks(cpu_id_p1) = n_blocks(cpu_id_p1) + 1

      if (.not. current%var_allocate_flag) then
         allocate(current%ro(nx_block, ny_block), current%mx(nx_block, ny_block), current%my(nx_block, ny_block),                  &
                  current%mz(nx_block, ny_block), current%bx(nx_block, ny_block), current%by(nx_block, ny_block),                  &
                  current%bz(nx_block, ny_block), current%en(nx_block, ny_block), current%gx(nx_block, ny_block),                  &
                  current%gy(nx_block, ny_block), current%po(nx_block, ny_block), current%az(nx_block, ny_block))
         allocate(current%fxro_bnd(ny_block, n_neighs), current%fxmx_bnd(ny_block, n_neighs),                                      &
                  current%fxmy_bnd(ny_block, n_neighs), current%fxmz_bnd(ny_block, n_neighs),                                      &
                  current%fxbx_bnd(ny_block, n_neighs), current%fxby_bnd(ny_block, n_neighs),                                      &
                  current%fxbz_bnd(ny_block, n_neighs), current%fxen_bnd(ny_block, n_neighs),                                      &
                  current%fyro_bnd(nx_block, n_neighs), current%fymx_bnd(nx_block, n_neighs),                                      &
                  current%fymy_bnd(nx_block, n_neighs), current%fymz_bnd(nx_block, n_neighs),                                      &
                  current%fybx_bnd(nx_block, n_neighs), current%fyby_bnd(nx_block, n_neighs),                                      &
                  current%fybz_bnd(nx_block, n_neighs), current%fyen_bnd(nx_block, n_neighs),                                      &
                  current%fxpo_bnd(ny_block, n_neighs), current%fypo_bnd(nx_block, n_neighs))
         current%var_allocate_flag = .true.
         if (ini_value_save_flag .ne. 0) then
            allocate(current%ro0(nx_block, ny_block), current%mx0(nx_block, ny_block), current%my0(nx_block, ny_block),            &
                     current%mz0(nx_block, ny_block), current%bx0(nx_block, ny_block), current%by0(nx_block, ny_block),            &
                     current%bz0(nx_block, ny_block), current%en0(nx_block, ny_block), current%az0(nx_block, ny_block))
         endif
      endif

      if (ini_value_save_flag .ne. 0) then
         do j = 1, ny_block
         do i = 1, nx_block
            current%ro(i, j) = info_var_recv(1, i, j, i_hilbert)
            current%mx(i, j) = info_var_recv(2, i, j, i_hilbert)
            current%my(i, j) = info_var_recv(3, i, j, i_hilbert)
            current%mz(i, j) = info_var_recv(4, i, j, i_hilbert)
            current%bx(i, j) = info_var_recv(5, i, j, i_hilbert)
            current%by(i, j) = info_var_recv(6, i, j, i_hilbert)
            current%bz(i, j) = info_var_recv(7, i, j, i_hilbert)
            current%en(i, j) = info_var_recv(8, i, j, i_hilbert)
            current%az(i, j) = info_var_recv(9, i, j, i_hilbert)
            current%gx(i, j) = info_var_recv(10, i, j, i_hilbert)
            current%gy(i, j) = info_var_recv(11, i, j, i_hilbert)
            current%po(i, j) = info_var_recv(12, i, j, i_hilbert)
            current%ro0(i, j) = info_var_recv(13, i, j, i_hilbert)
            current%mx0(i, j) = info_var_recv(14, i, j, i_hilbert)
            current%my0(i, j) = info_var_recv(15, i, j, i_hilbert)
            current%mz0(i, j) = info_var_recv(16, i, j, i_hilbert)
            current%bx0(i, j) = info_var_recv(17, i, j, i_hilbert)
            current%by0(i, j) = info_var_recv(18, i, j, i_hilbert)
            current%bz0(i, j) = info_var_recv(19, i, j, i_hilbert)
            current%en0(i, j) = info_var_recv(20, i, j, i_hilbert)
            current%az0(i, j) = info_var_recv(21, i, j, i_hilbert)
         enddo
         enddo
      else
         do j = 1, ny_block
         do i = 1, nx_block
            current%ro(i, j) = info_var_recv(1, i, j, i_hilbert)
            current%mx(i, j) = info_var_recv(2, i, j, i_hilbert)
            current%my(i, j) = info_var_recv(3, i, j, i_hilbert)
            current%mz(i, j) = info_var_recv(4, i, j, i_hilbert)
            current%bx(i, j) = info_var_recv(5, i, j, i_hilbert)
            current%by(i, j) = info_var_recv(6, i, j, i_hilbert)
            current%bz(i, j) = info_var_recv(7, i, j, i_hilbert)
            current%en(i, j) = info_var_recv(8, i, j, i_hilbert)
            current%az(i, j) = info_var_recv(9, i, j, i_hilbert)
            current%gx(i, j) = info_var_recv(10, i, j, i_hilbert)
            current%gy(i, j) = info_var_recv(11, i, j, i_hilbert)
            current%po(i, j) = info_var_recv(12, i, j, i_hilbert)
         enddo
         enddo
      endif

   enddo

!----------------------------------------------------------------------------------------------------------------------------------|
!  update the links
!----------------------------------------------------------------------------------------------------------------------------------|
   hilbert_start = hilbert_start_new
   hilbert_end = hilbert_end_new

   last => base
   current => base%framework_next
   do while (current%lv .ne. 0)
      if (current%hilbert_number .ge. hilbert_start(cpu_id_p1) .and. current%hilbert_number .le. hilbert_end(cpu_id_p1)) then
         max_level = max(max_level, current%lv)
         current%last => last
         last%next => current
         last => current
      else
         current%last => null()
         current%next => null()
         if (current%var_allocate_flag) then
            deallocate(current%ro, current%mx, current%my, current%mz, current%bx, current%by, current%bz, current%en,             &
                       current%gx, current%gy, current%po, current%az)
            deallocate(current%fxro_bnd, current%fxmx_bnd, current%fxmy_bnd, current%fxmz_bnd,                                     &
                       current%fxbx_bnd, current%fxby_bnd, current%fxbz_bnd, current%fxen_bnd,                                     &
                       current%fyro_bnd, current%fymx_bnd, current%fymy_bnd, current%fymz_bnd,                                     &
                       current%fybx_bnd, current%fyby_bnd, current%fybz_bnd, current%fyen_bnd,                                     &
                       current%fxpo_bnd, current%fypo_bnd)
            current%var_allocate_flag = .false.
            if (ini_value_save_flag .ne. 0) then
               deallocate(current%ro0, current%mx0, current%my0, current%mz0, current%bx0, current%by0, current%bz0, current%en0,  &
                  current%az0)
            endif
         endif
      endif
      current => current%framework_next
   enddo
   base%last => last
   last%next => base

!----------------------------------------------------------------------------------------------------------------------------------|
!  calculation the total number of the blocks and update the n_blocks
!----------------------------------------------------------------------------------------------------------------------------------|
   call mpi_allreduce (n_blocks(cpu_id_p1), n_blocks_all, 1, mpi_integer4, mpi_sum, mpi_comm_world, error)
   average_blocks = n_blocks(cpu_id_p1)
   call mpi_allgather (average_blocks, 1, mpi_integer4, n_blocks, 1, mpi_integer4, mpi_comm_world, error)

!----------------------------------------------------------------------------------------------------------------------------------|
!  free memory
!----------------------------------------------------------------------------------------------------------------------------------|
   deallocate (request, request_status, info_int_send, info_var_send, info_int_recv, info_var_recv)

   return
end subroutine load_balance
